package com.zjj.lbw.ai;

import dev.langchain4j.data.message.ChatMessage;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.memory.chat.TokenWindowChatMemory;
import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.model.openai.OpenAiTokenizer;
import dev.langchain4j.model.zhipu.ZhipuAiChatModel;
import dev.langchain4j.service.AiServices;
import dev.langchain4j.service.MemoryId;
import dev.langchain4j.service.UserMessage;
import dev.langchain4j.store.memory.chat.ChatMemoryStore;
import org.mapdb.DB;
import org.mapdb.DBMaker;

import java.util.List;
import java.util.Map;

import static dev.langchain4j.data.message.ChatMessageDeserializer.messagesFromJson;
import static dev.langchain4j.data.message.ChatMessageSerializer.messagesToJson;
import static org.mapdb.Serializer.STRING;

public class NamingMain {
    interface NameMaster {
        String name(@MemoryId String userId, @UserMessage String desc);
    }

    public static void main(String[] args) {
        ChatLanguageModel model = ZhipuAiChatModel.builder()
                .apiKey("9a6dfc4258104e28ad0c72db2ddda3e4.f6VL6AWRQxORuJpi")
                .build();

//        MessageWindowChatMemory memory = MessageWindowChatMemory.withMaxMessages(10);
//        NameMaster nameMaster = AiServices.builder(NameMaster.class)
//                .chatLanguageModel(model)
//                .chatMemory(memory)
//                .build();

//        MessageWindowChatMemory memory = MessageWindowChatMemory.withMaxMessages(10);
//        NameMaster nameMaster = AiServices.builder(NameMaster.class)
//                .chatLanguageModel(model)
//                .chatMemoryProvider(memoryId -> MessageWindowChatMemory.withMaxMessages(10))
//                .build();

        NameMaster nameMaster = AiServices.builder(NameMaster.class)
                .chatLanguageModel(model)
                .chatMemoryProvider(userId -> TokenWindowChatMemory.builder()
                        .id(userId)
                        .maxTokens(10000, new OpenAiTokenizer())
                        .chatMemoryStore(new PersistentChatMemoryStore())
                        .build())
                .build();

//        String name = nameMaster.name("1", "帮我取一个很有中国文化内涵的男孩的名字，给我一个你觉得最好的");
//        System.out.println(name);

        String name1 = nameMaster.name("1", "换一个");
        System.out.println(name1);
    }

    static class PersistentChatMemoryStore implements ChatMemoryStore {

        private final DB db = DBMaker.fileDB("chat-memory.db").transactionEnable().make();
        private final Map<String, String> map = db.hashMap("messages", STRING, STRING).createOrOpen();

        @Override
        public List<ChatMessage> getMessages(Object memoryId) {
            String json = map.get((String) memoryId);
            return messagesFromJson(json);
        }

        @Override
        public void updateMessages(Object memoryId, List<ChatMessage> messages) {
            String json = messagesToJson(messages);
            map.put((String) memoryId, json);
            db.commit();
        }

        @Override
        public void deleteMessages(Object memoryId) {
            map.remove((String) memoryId);
            db.commit();
        }
    }
}
